from typing import Optional
from typing_extensions import Literal

import io

import numpy as np
from PIL import Image


def nda_to_img(
    nda_data: np.ndarray,
    nda_mask: Optional[np.ndarray] = None,
    dim_order: Literal["CHW", "HWC"] = "HWC",
    format: Literal["PNG", "WebP", "JPEG"] = "WebP",
    fpath_dest: Optional[str] = None,
) -> Optional[io.BytesIO]:
    """Convert a numpy array to a PIL image

    Parameters
    ----------
    nda_data :
        A numpy array in the shape of [n_rows X n_cols] or [n_rows X n_cols X n_bands] to be converted.
        If with multiple bands expected to be 3 (RGB) or 4 (RGBA) bands.
        Notice it is the callers' duty to make sure valid values are provided in `nda_data`,
        for floating dtype, values should be within the range [0, 1],
        if uint8 dtype, values should be within the range [0, 255].
    nda_mask : Optional[np.ndarray], optional
        A numpy array with uint8 or bool dtype for the mask,
        with valid pixels of value 1 and nonvalid pixels of value 0,
        by default None.
    dim_order : Literal["CHW", "HWC"], optional
        The order of channel/band, by default "HWC" putting band in the last.
    format : Literal["PNG", "WebP", "JPEG"], optional
        Format to use, by default "WebP", other options are "PNG" and "JPEG".
    fpath_dest : Optional[str], optional
        Destination file path, by default None not to write to file but to return a BytesIO object of the image.

    Returns
    -------
    Optional[io.BytesIO]
        a BytesIO object of the image if `fpath_dest` is None, otherwise None and the result would be written to file.
    """
    if dim_order == "CHW":
        # Convert CxHxW to HxWxC
        nda_data = np.moveaxis(nda_data, 0, -1)

    # nda_data could be a single band or 3 bands (band first) ndarray
    im_to_show = Image.fromarray(nda_data)
    if nda_mask is not None:
        im_to_show.putalpha(Image.fromarray(nda_mask))

    if fpath_dest is not None:
        im_to_show.save(fpath_dest, format=format)
    else:
        img_byte_arr = io.BytesIO()
        im_to_show.save(img_byte_arr, format=format)
        img_byte_arr.seek(0)
        return img_byte_arr
